import pygame
import math

# 初始化 Pygame
pygame.init()

# 定义常量
WIDTH, HEIGHT = 800, 600
TABLE_COLOR = (0, 128, 0)
BALL_COLOR = (255, 255, 255)
CUE_COLOR = (128, 64, 0)
BALL_RADIUS = 10
CUE_LENGTH = 100
CUE_WIDTH = 5
FRICTION = 0.99

# 创建游戏窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("桌球游戏")

# 球的类
class Ball:
    def __init__(self, x, y, radius, color):
        self.x = x
        self.y = y
        self.radius = radius
        self.color = color
        self.vx = 0
        self.vy = 0

    def draw(self):
        pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), self.radius)

    def move(self):
        self.x += self.vx
        self.y += self.vy
        self.vx *= FRICTION
        self.vy *= FRICTION
        # 边界检测
        if self.x - self.radius < 0 or self.x + self.radius > WIDTH:
            self.vx = -self.vx
        if self.y - self.radius < 0 or self.y + self.radius > HEIGHT:
            self.vy = -self.vy

    def is_colliding(self, other):
        distance = math.sqrt((self.x - other.x) ** 2 + (self.y - other.y) ** 2)
        return distance <= self.radius + other.radius

    def collide(self, other):
        dx = other.x - self.x
        dy = other.y - self.y
        distance = math.sqrt(dx ** 2 + dy ** 2)
        normal_x = dx / distance
        normal_y = dy / distance
        v1n = self.vx * normal_x + self.vy * normal_y
        v2n = other.vx * normal_x + other.vy * normal_y
        v1t = self.vx - v1n * normal_x
        v2t = other.vx - v2n * normal_x
        v1n_new = v2n
        v2n_new = v1n
        self.vx = v1t + v1n_new * normal_x
        self.vy = v1t + v1n_new * normal_y
        other.vx = v2t + v2n_new * normal_x
        other.vy = v2t + v2n_new * normal_y

# 创建球
cue_ball = Ball(WIDTH // 2, HEIGHT // 2, BALL_RADIUS, BALL_COLOR)
target_ball = Ball(WIDTH // 2 + 100, HEIGHT // 2, BALL_RADIUS, BALL_COLOR)

# 球杆状态
cue_angle = 0
cue_power = 0
shooting = False

# 游戏主循环
clock = pygame.time.Clock()
running = True
while running:
    screen.fill(TABLE_COLOR)

    # 处理事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.MOUSEBUTTONDOWN:
            shooting = True
        elif event.type == pygame.MOUSEBUTTONUP:
            shooting = False
            cue_ball.vx = cue_power * math.cos(cue_angle)
            cue_ball.vy = cue_power * math.sin(cue_angle)
            cue_power = 0

    if shooting:
        mouse_x, mouse_y = pygame.mouse.get_pos()
        dx = mouse_x - cue_ball.x
        dy = mouse_y - cue_ball.y
        cue_angle = math.atan2(dy, dx)
        cue_power = math.sqrt(dx ** 2 + dy ** 2) / 10

    # 绘制球杆
    if shooting:
        cue_end_x = cue_ball.x - CUE_LENGTH * math.cos(cue_angle)
        cue_end_y = cue_ball.y - CUE_LENGTH * math.sin(cue_angle)
        pygame.draw.line(screen, CUE_COLOR, (cue_ball.x, cue_ball.y), (cue_end_x, cue_end_y), CUE_WIDTH)

    # 球的移动和碰撞检测
    cue_ball.move()
    target_ball.move()
    if cue_ball.is_colliding(target_ball):
        cue_ball.collide(target_ball)

    # 绘制球
    cue_ball.draw()
    target_ball.draw()

    # 更新显示
    pygame.display.flip()
    clock.tick(60)

# 退出游戏
pygame.quit()